import torch
import torch.nn as nn
import numpy as np


class RMatrix(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, triangles, barycenters, indices_neigh_tri, number_neigh_tri):
        # 计算重心差异
        # barycenters_first: 获取每个邻接三角形对中的第一个三角形的重心
        # barycenters_second: 获取每个邻接三角形对中的第二个三角形的重心
        barycenters_first = barycenters[indices_neigh_tri[:, 0]].unsqueeze(1)  # 添加一个新的维度
        barycenters_second = barycenters[indices_neigh_tri[:, 1:]]
        barycenters_diff = barycenters_first - barycenters_second  # 计算重心差异

        # 计算三角形边的长度（范数）
        # v0, v1, v2: 三角形的三个顶点
        v0, v1, v2 = triangles[:, 0], triangles[:, 1], triangles[:, 2]

        # 计算三角形的边长向量
        e_ij = torch.norm(v0 - v1, dim=1)  # 边e_ij的长度
        e_ik = torch.norm(v0 - v2, dim=1)  # 边e_ik的长度
        e_jk = torch.norm(v1 - v2, dim=1)  # 边e_jk的长度

        # 将三条边的长度沿着最后一个维度堆叠起来，得到一个(N, 3)的张量
        diff_vectors = torch.stack([e_ij, e_ik, e_jk], dim=1)

        # 计算最大值和最小值的边长差异
        max_diff_vectors = diff_vectors.max(dim=1).values  # 计算每个三角形的最大边长
        min_diff_vectors = diff_vectors.min(dim=1).values  # 计算每个三角形的最小边长

        # 计算最大边长差异
        max_diff_vectors_diff = max_diff_vectors[indices_neigh_tri[:, 0]][:, None] - max_diff_vectors[
            indices_neigh_tri[:, 1:]]
        # 计算最小边长差异
        min_diff_vectors_diff = min_diff_vectors[indices_neigh_tri[:, 0]][:, None] - min_diff_vectors[
            indices_neigh_tri[:, 1:]]

        # 计算R矩阵，将最小边长差异、最大边长差异和重心差异拼接在一起
        r_matrix = torch.cat((min_diff_vectors_diff.unsqueeze(2), max_diff_vectors_diff.unsqueeze(2), barycenters_diff),
                             dim=2)

        return r_matrix
